﻿/*
 * @author: wizardc
 */

using System.Collections.Generic;
using System.IO;
using System.Linq;
using Dou.BoundingVolumeHierarchy;
using Dou.Cluster;
using DouEditor.Utils;
using Unity.Mathematics;
using UnityEditor;
using UnityEngine;

namespace Demo.WorldMap
{
    /// <summary>
    /// 场景静态渲染数据导出类
    /// 1. 预制类型的数据量是较少的（一般用于如雾，入海波浪等点缀特效或动画之类），所以采用简单直观的方式来直接导出BVH数据即可；
    /// 2. GPU绘制类型的数据量就比较庞大了（一般用于树木植被山脉等大量出现的装饰），所以用上了蔟划分即使每个叶子节点仅可能包含多个
    /// 附近的资源（避免BVH层数过多）、数据流（将大量的转换信息存储为流文件）、同时为了避免长时间构建还对BVH的构建进行了拆分合并的
    /// 额外操作；
    /// </summary>
    public static class SceneStaticDataExport
    {

        #region 构建简单 BVH

        /// <summary>
        /// 构建为简单 BVH
        /// </summary>
        public static void BuildSimpleBVH(Dictionary<string, List<GameObject>> gameObjectMap, out List<PrefabInfo> prefabInfoList, out Dictionary<string, BVH<SimpleBVHLeaf>> binaryTreeMap)
        {
            prefabInfoList = new List<PrefabInfo>();
            binaryTreeMap = new Dictionary<string, BVH<SimpleBVHLeaf>>();
            
            var assetIndexMap = new Dictionary<GameObject, int>();
            foreach (var kv in gameObjectMap)
            {
                var name = kv.Key;
                var list = kv.Value;
                if (list.Count == 0)
                {
                    continue;
                }

                var bvh = BuildSimpleBVH(prefabInfoList, list, assetIndexMap);
                binaryTreeMap.Add(name, bvh);
            }
        }

        private static BVH<SimpleBVHLeaf> BuildSimpleBVH(List<PrefabInfo> prefabInfoList, List<GameObject> gameObjectList, Dictionary<GameObject, int> assetIndexMap)
        {
            var bvh = new BVH<SimpleBVHLeaf>();

            foreach (var gameObject in gameObjectList)
            {
                var instanceObj = new PrefabInstanceObject
                {
                    prefabIndex = GetPrefabIndex(gameObject, prefabInfoList, assetIndexMap),
                    bounds = EditorUtil.GetBounds(gameObject),
                    position = gameObject.transform.position,
                    rotation = gameObject.transform.rotation,
                    scale = gameObject.transform.localScale
                };
                
                var leaf = new SimpleBVHLeaf
                {
                    instance = instanceObj
                };
                bvh.Add(leaf);
            }
            
            bvh.Rebuild();

            return bvh;
        }
        
        private static int GetPrefabIndex(GameObject prefab, List<PrefabInfo> prefabInfoList, Dictionary<GameObject, int> assetIndexMap)
        {
            var asset = AssetDatabase.LoadAssetAtPath<GameObject>(PrefabUtility.GetPrefabAssetPathOfNearestInstanceRoot(prefab));
            if (assetIndexMap.ContainsKey(asset))
            {
                return assetIndexMap[asset];
            }

            var info = new PrefabInfo
            {
                asset = asset,
                bounds = EditorUtil.GetBounds(asset)
            };
            prefabInfoList.Add(info);
            var index = prefabInfoList.Count - 1;
            assetIndexMap.Add(asset, index);
            return index;
        }
        
        #endregion
        
        #region 导出预制类型配置

        public static PrefabStaticRenderData ExportPrefabConfig(List<PrefabInfo> prefabInfoList, Dictionary<string, BVH<SimpleBVHLeaf>> bvhMap)
        {
            var config = ScriptableObject.CreateInstance<PrefabStaticRenderData>();

            var resList = new PrefabStaticRenderRes[prefabInfoList.Count];
            for (int i = 0; i < prefabInfoList.Count; i++)
            {
                var prefabInfo = prefabInfoList[i];
                var res = new PrefabStaticRenderRes();
#if UNITY_EDITOR
                res.prefab = prefabInfo.asset;
#endif
                res.prefabName = Path.GetFileNameWithoutExtension(AssetDatabase.GetAssetPath(prefabInfo.asset));
                res.bounds = prefabInfo.bounds;
                resList[i] = res;
            }
            config.res = resList;

            var trees = new PrefabStaticRenderTree[bvhMap.Count];
            var index = 0;
            foreach (var kv in bvhMap)
            {
                var name = kv.Key;
                var tree = kv.Value;
                
                var renderTree = new PrefabStaticRenderTree();
                renderTree.name = name;
                
                var nodes = new PrefabStaticRenderNode[tree.nodeCount];
                var matrixes = new List<Matrix4x4>();
                WritePrefabBVHNode(tree, nodes, matrixes);
                renderTree.nodes = nodes;
                renderTree.matrixes = matrixes.ToArray();
                
                trees[index++] = renderTree;
            }
            config.trees = trees;
            
            return config;
        }

        private static void WritePrefabBVHNode(BVH<SimpleBVHLeaf> tree, PrefabStaticRenderNode[] nodes, List<Matrix4x4> matrixes)
        {
            var queue = new Queue<int>();
            queue.Enqueue(tree.rootID);
            
            // 当前写入的索引
            var currentIndex = 0;
            // 插入子节点的索引
            var insertIndex = 0;
            
            while (queue.Count > 0)
            {
                var treeIndex = queue.Dequeue();
                var treeNode = tree.GetNode(treeIndex);
                
                var node = new PrefabStaticRenderNode();
                
                node.index = currentIndex;
                // 子节点索引
                node.child1 = treeNode.isLeaf ? -1 : ++insertIndex;
                node.child2 = treeNode.isLeaf ? -1 : ++insertIndex;
                
                node.bounds = treeNode.aabb;
                
                if (treeNode.isLeaf)
                {
                    var instance = treeNode.obj.instance;
                    node.resId = instance.prefabIndex;
                    node.matrixId = matrixes.Count;
                    var mat = Matrix4x4.TRS(instance.position, instance.rotation, instance.scale);
                    matrixes.Add(mat);
                }
                else
                {
                    node.resId = -1;
                    node.matrixId = -1;
                }
                
                nodes[currentIndex++] = node;
                
                if (!treeNode.isLeaf)
                {
                    queue.Enqueue(treeNode.child1);
                    queue.Enqueue(treeNode.child2);
                }
            }
        }

        #endregion

        #region 构建蔟 BVH

        /// <summary>
        /// 构建为蔟 BVH
        /// </summary>
        /// <param name="maxSize">一个蔟最大的尺寸限制</param>
        /// <param name="maxCount">一个蔟里最多的数量限制</param>
        /// <param name="splitChunkX">分块生成BVH时的X轴长度，小于等于0时表示不分块，如果对象数量过大生成BVH会非常卡顿，通过该参数拆分为多个BVH生成加快导出效率</param>
        /// <param name="splitChunkZ">分块生成BVH时的Z轴长度，小于等于0时表示不分块，如果对象数量过大生成BVH会非常卡顿，通过该参数拆分为多个BVH生成加快导出效率</param>
        public static void BuildClusterBVH(Dictionary<string, List<GameObject>> gameObjectMap, float maxSize, int maxCount, float splitChunkX, float splitChunkZ, out List<PrefabInfo> prefabInfoList, out Dictionary<string, ClusterBinaryTree> binaryTreeMap)
        {
            prefabInfoList = new List<PrefabInfo>();
            binaryTreeMap = new Dictionary<string, ClusterBinaryTree>();

            var assetIndexMap = new Dictionary<GameObject, int>();
            foreach (var kv in gameObjectMap)
            {
                var name = kv.Key;
                var list = kv.Value;
                if (list.Count == 0)
                {
                    continue;
                }

                var clusterList = ClusterSeparate(list, maxSize, maxCount);

                // // 调试
                // var count = 0;
                // foreach (var item in clusterList)
                // {
                //     count += item.Count;
                // }
                // Debug.Log($"蔟拆分后的数据量: {count}，拆分前的数据量：{list.Count}");
                
                var leaves = GetLeaves(clusterList, prefabInfoList, assetIndexMap);
                var bvh = BuildBVH(leaves, splitChunkX, splitChunkZ);
                binaryTreeMap.Add(name, bvh);
            }
        }
        
        /// <summary>
        /// 使用聚类算法进行分组划分
        /// </summary>
        private static List<List<GameObject>> ClusterSeparate(List<GameObject> gameObjectList, float maxSize, int maxCount)
        {
            var list = new List<GameObjectCluster>();
            foreach (var gameObject in gameObjectList)
            {
                list.Add(new GameObjectCluster
                {
                    gameObject = gameObject
                });
            }

            var separates = new List<List<GameObjectCluster>>();
            ClusterSeparate(list, maxSize, maxCount, separates);
            
            var result = new List<List<GameObject>>();
            foreach (var separate in separates)
            {
                var goList = new List<GameObject>();
                foreach (var cluster in separate)
                {
                    goList.Add(cluster.gameObject);
                }
                result.Add(goList);
            }
            return result;
        }

        /// <summary>
        /// 使用聚类算法进行分组划分
        /// </summary>
        private static void ClusterSeparate(List<GameObjectCluster> list, float maxSize, int maxCount, List<List<GameObjectCluster>> result)
        {
            var bounds = GetBounds(list);
            if (list.Count >= 2 && Max(bounds.size) > maxSize)
            {
                var cluster = new BinarySearchKMeansCluster<GameObjectCluster>(list);
                cluster.Cluster(list);
                var clusterList = cluster.clusterList;
                var firstCluster = clusterList[0];
                var secondCluster = clusterList[1];
                var firstBounds = GetBounds(firstCluster);
                var secondBounds = GetBounds(secondCluster);
                // 处理第一个聚落
                if (firstCluster.Count >= 2 && (Max(firstBounds.size) > maxSize || firstCluster.Count > maxCount))
                {
                    ClusterSeparate(firstCluster, maxSize, maxCount, result);
                }
                else if (firstCluster.Count > 0)
                {
                    result.Add(firstCluster);
                }
                // 处理第二个聚落
                if (secondCluster.Count >= 2 && (Max(secondBounds.size) > maxSize || secondCluster.Count > maxCount))
                {
                    ClusterSeparate(secondCluster, maxSize, maxCount, result);
                }
                else if (secondCluster.Count > 0)
                {
                    result.Add(secondCluster);
                }
            }
            else if (list.Count > 0)
            {
                result.Add(list);
            }
        }

        private static Bounds GetBounds(List<GameObjectCluster> list)
        {
            var bounds = new Bounds();
            var isFirst = true;
            foreach (var cluster in list)
            {
                var b = EditorUtil.GetBounds(cluster.gameObject);
                if (isFirst)
                {
                    bounds = b;
                    isFirst = false;
                }
                else
                {
                    bounds.Encapsulate(b);
                }
            }
            return bounds;
        }
        
        private static float Max(Vector3 v)
        {
            return Mathf.Max(v.x, v.y, v.z);
        }
        
        /// <summary>
        /// 将所有的数据转换为BVH的叶子节点
        /// </summary>
        private static List<ClusterBVHLeaf> GetLeaves(List<List<GameObject>> clusterList, List<PrefabInfo> prefabInfoList, Dictionary<GameObject, int> assetIndexMap)
        {
            var leafList = new List<ClusterBVHLeaf>();
            foreach (var cluster in clusterList)
            {
                var node = new ClusterNode();
                var bounds = new Bounds();
                var isFirst = true;
                node.instanceList = new List<PrefabInstanceObject>();
                foreach (var gameObject in cluster)
                {
                    var instanceObj = new PrefabInstanceObject
                    {
                        prefabIndex = GetPrefabIndex(gameObject, prefabInfoList, assetIndexMap),
                        bounds = EditorUtil.GetBounds(gameObject),
                        position = gameObject.transform.position,
                        rotation = gameObject.transform.rotation,
                        scale = gameObject.transform.localScale
                    };
                    node.instanceList.Add(instanceObj);
                    if (isFirst)
                    {
                        bounds = instanceObj.bounds;
                        isFirst = false;
                    }
                    else
                    {
                        bounds.Encapsulate(instanceObj.bounds);
                    }
                }
                node.bounds = bounds;
                
                var leaf = new ClusterBVHLeaf
                {
                    node = node
                };
                leafList.Add(leaf);
            }
            return leafList;
        }

        /// <summary>
        /// 构建为 BVH 树
        /// </summary>
        private static ClusterBinaryTree BuildBVH(List<ClusterBVHLeaf> leaves, float splitChunkX, float splitChunkZ)
        {
            // 不需要拆分构建的情况
            if (splitChunkX <= 0 && splitChunkZ <= 0)
            {
                var bvh = new BVH<ClusterBVHLeaf>();
                foreach (var leaf in leaves)
                {
                    bvh.Add(leaf);
                }
                bvh.Rebuild();
                return Bvh2BinaryTree(bvh);
            }

            // 拆分为多个小块构建 BVH 后再合并
            var bounds = GetBounds(leaves);
            var col = Mathf.CeilToInt(bounds.size.x / splitChunkX);
            var row = Mathf.CeilToInt(bounds.size.z / splitChunkZ);
            
            var total = col * row;
            var bvhArr = new BVH<ClusterBVHLeaf>[total];
            for (var i = 0; i < total; i++)
            {
                bvhArr[i] = new BVH<ClusterBVHLeaf>();
            }
            
            foreach (var leaf in leaves)
            {
                var pos = leaf.GetPosition();
                var c = Mathf.FloorToInt(pos.x / splitChunkX);
                var r = Mathf.FloorToInt(pos.z / splitChunkZ);
                var index = r * col + c;
                var bvh = bvhArr[index];
                bvh.Add(leaf);
            }
            
            for (var i = 0; i < total; i++)
            {
                var bvh = bvhArr[i];
                if (!bvh.rootIsNull)
                {
                    bvh.Rebuild();
                }
            }
            
            return MergeBVH(bvhArr);
        }

        /// <summary>
        /// 将 BVH 转换为自定义的二叉树数据，方便对数进行合并操作
        /// </summary>
        private static ClusterBinaryTree Bvh2BinaryTree(BVH<ClusterBVHLeaf> bvh)
        {
            var trees = new ClusterBinaryTreeNode[bvh.nodeCount];
            
            var queue = new Queue<int>();
            queue.Enqueue(bvh.rootID);
            while (queue.Count > 0)
            {
                var bvhId = queue.Dequeue();
                var node = bvh.GetNode(bvhId);
                
                var btNode = new ClusterBinaryTreeNode
                {
                    child1Index = node.child1,
                    child2Index = node.child2,
                    aabb = node.aabb
                };
                if (node.isLeaf)
                {
                    btNode.obj = node.obj.node.instanceList;
                    btNode.subBvhIndex = node.obj.node.subBvhIndex;
                }
                trees[bvhId] = btNode;

                if (!node.isLeaf)
                {
                    queue.Enqueue(node.child1);
                    queue.Enqueue(node.child2);
                }
            }

            var tree = new ClusterBinaryTree
            {
                root = trees[bvh.rootID],
                nodeCount = bvh.nodeCount
            };
            for (var i = 0; i < trees.Length; i++)
            {
                var node = trees[i];
                if (!node.isLeaf)
                {
                    node.child1 = trees[node.child1Index];
                    node.child2 = trees[node.child2Index];
                }
            }
            return tree;
        }
        
        private static Bounds GetBounds(List<ClusterBVHLeaf> list)
        {
            var bounds = new Bounds();
            var isFirst = true;
            foreach (var leaf in list)
            {
                var b = leaf.GetBounds();
                if (isFirst)
                {
                    bounds = b;
                    isFirst = false;
                }
                else
                {
                    bounds.Encapsulate(b);
                }
            }
            return bounds;
        }

        /// <summary>
        /// 合并多个 BVH 树为一个
        /// </summary>
        private static ClusterBinaryTree MergeBVH(BVH<ClusterBVHLeaf>[] bvhArr)
        {
            // 用每个BVH的根节点来构建根 BVH
            var rootBvh = new BVH<ClusterBVHLeaf>();
            for (var i = 0; i < bvhArr.Length; i++)
            {
                var bvh = bvhArr[i];
                if (!bvh.rootIsNull)
                {
                    var leaf = new ClusterBVHLeaf
                    {
                        node = new ClusterNode
                        {
                            bounds = bvh.GetNode(bvh.rootID).aabb,
                            subBvhIndex = i // 记录索引，合并时关联使用
                        }
                    };
                    rootBvh.Add(leaf);
                }
            }
            rootBvh.Rebuild();
            
            // 合并 BVH
            var root = Bvh2BinaryTree(rootBvh);
            var queue = new Queue<ClusterBinaryTreeNode>();
            queue.Enqueue(root.root);
            while (queue.Count > 0)
            {
                var node = queue.Dequeue();
                if (!node.isLeaf)
                {
                    if (node.child1.isLeaf)
                    {
                        var rootChild = Bvh2BinaryTree(bvhArr[node.child1.subBvhIndex]);
                        node.child1 = rootChild.root;
                        root.nodeCount += rootChild.nodeCount - 1;
                    }
                    else
                    {
                        queue.Enqueue(node.child1);
                    }
                    if (node.child2.isLeaf)
                    {
                        var rootChild = Bvh2BinaryTree(bvhArr[node.child2.subBvhIndex]);
                        node.child2 = rootChild.root;
                        root.nodeCount += rootChild.nodeCount - 1;
                    }
                    else
                    {
                        queue.Enqueue(node.child2);
                    }
                }
            }
            return root;
        }

        #endregion

        #region 导出GPU类型配置

        public static GPUStaticRenderData ExportGPUConfig(List<PrefabInfo> prefabInfoList, Dictionary<string, ClusterBinaryTree> binaryTreeMap, List<float4x4> transList)
        {
            var config = ScriptableObject.CreateInstance<GPUStaticRenderData>();

            var resList = new List<GPUStaticRenderRes>();
            var lodList = new List<GPUStaticRenderLODRes>();
            var resMap = new Dictionary<int, int>();
            var lodResMap = new Dictionary<string, int>();
            
            var trees = new GPUStaticRenderTree[binaryTreeMap.Count];
            var index = 0;
            foreach (var kv in binaryTreeMap)
            {
                var name = kv.Key;
                var tree = kv.Value;

                resList.Clear();
                lodList.Clear();
                resMap.Clear();
                lodResMap.Clear();
                
                var renderTree = new GPUStaticRenderTree();
                renderTree.name = name;
                
                var nodes = new GPUStaticRenderNode[tree.nodeCount];
                var leaves = new List<GPUStaticRenderLeaf>();
                WriteGPUBVHNode(prefabInfoList, tree.root, nodes, leaves, transList, resList, lodList, resMap, lodResMap);
                renderTree.res = resList.ToArray();
                renderTree.lodRes = lodList.ToArray();
                renderTree.nodes = nodes;
                renderTree.leaves = leaves.ToArray();
                
                trees[index++] = renderTree;
            }
            config.trees = trees;

            return config;
        }

        private static void WriteGPUBVHNode(List<PrefabInfo> prefabInfoList, ClusterBinaryTreeNode treeRoot, GPUStaticRenderNode[] nodes, List<GPUStaticRenderLeaf> leafList, List<float4x4> transList, List<GPUStaticRenderRes> resList, List<GPUStaticRenderLODRes> lodList, Dictionary<int, int> resMap, Dictionary<string, int> lodResMap)
        {
            var queue = new Queue<ClusterBinaryTreeNode>();
            queue.Enqueue(treeRoot);
            
            // 当前写入的索引
            var currentIndex = 0;
            // 插入子节点的索引
            var insertIndex = 0;

            while (queue.Count > 0)
            {
                var treeNode = queue.Dequeue();
                
                var node = new GPUStaticRenderNode();
                
                node.index = currentIndex;
                // 子节点索引
                node.child1 = treeNode.isLeaf ? -1 : ++insertIndex;
                node.child2 = treeNode.isLeaf ? -1 : ++insertIndex;

                node.bounds = treeNode.aabb;
                
                if (treeNode.isLeaf)
                {
                    node.startLeafIndex = leafList.Count;
                    node.leafCount = WriteGPUBVHLeaf(prefabInfoList, treeNode.obj, leafList, transList, resList, lodList, resMap, lodResMap);
                }
                else
                {
                    node.startLeafIndex = -1;
                    node.leafCount = 0;
                }
                
                nodes[currentIndex++] = node;
                
                if (!treeNode.isLeaf)
                {
                    queue.Enqueue(treeNode.child1);
                    queue.Enqueue(treeNode.child2);
                }
            }
        }

        private static int WriteGPUBVHLeaf(List<PrefabInfo> prefabInfoList, List<PrefabInstanceObject> instanceList, List<GPUStaticRenderLeaf> leafList, List<float4x4> transList, List<GPUStaticRenderRes> resList, List<GPUStaticRenderLODRes> lodList, Dictionary<int, int> resMap, Dictionary<string, int> lodResMap)
        {
            instanceList.Sort((a, b) => a.prefabIndex.CompareTo(b.prefabIndex));
            var groups = instanceList.GroupBy(item => item.prefabIndex)
                .OrderBy(g => g.Key);
            var leafCount = groups.Count();
            foreach (var @group in groups)
            {
                var leaf = new GPUStaticRenderLeaf();
                leaf.resId = GetGPUResId(prefabInfoList, group.Key, resList, lodList, resMap, lodResMap);
                leaf.startDataIndex = transList.Count;
                leaf.dataCount = group.Count();
                foreach (var instanceObj in group)
                {
                    var trans = float4x4.TRS(instanceObj.position, instanceObj.rotation, instanceObj.scale);
                    trans.c0.w = math.max(instanceObj.bounds.extents.x, instanceObj.bounds.extents.z); // 存储当前对象的最大半径
                    trans.c1.w = instanceObj.bounds.extents.y; // 存储当前对象的高度的一半
                    transList.Add(trans);
                }
                leafList.Add(leaf);
            }
            return leafCount;
        }

        private static int GetGPUResId(List<PrefabInfo> prefabInfoList, int index, List<GPUStaticRenderRes> resList, List<GPUStaticRenderLODRes> lodList, Dictionary<int, int> resMap, Dictionary<string, int> lodResMap)
        {
            if (resMap.ContainsKey(index))
            {
                return resMap[index];
            }
            
            var res = new GPUStaticRenderRes();
            
            var prefabInfo = prefabInfoList[index];
            var lodGroup = prefabInfo.asset.GetComponent<LODGroup>();
            var lods = lodGroup.GetLODs();

            res.maxLod = lods.Length;
            res.lod = new GPUStaticRenderLODInfo[res.maxLod];
            for (int i = 0; i < lods.Length; i++)
            {
                var renderers = lods[i].renderers;
                var count = renderers.Length;

                res.lod[i] = new GPUStaticRenderLODInfo();
                res.lod[i].lodId = new int[count];
                var lodRes = res.lod[i].lodId;
                
                for (int j = 0; j < count; j++)
                {
                    var filter = renderers[j].GetComponent<MeshFilter>();
                    var mesh = filter.sharedMesh;
                    var mat = renderers[j].sharedMaterial;

                    lodRes[j] = GetGPULodResId(lodList, mesh, mat, lodResMap);
                }
            }
            res.bounds = prefabInfo.bounds;
            
            var resId = resList.Count;
            resList.Add(res);
            resMap.Add(index, resId);
            return resId;
        }

        private static int GetGPULodResId(List<GPUStaticRenderLODRes> lodList, Mesh mesh, Material mat, Dictionary<string, int> lodResMap)
        {
            var meshName = Path.GetFileNameWithoutExtension(AssetDatabase.GetAssetPath(mesh));
            var matName = Path.GetFileNameWithoutExtension(AssetDatabase.GetAssetPath(mat));

            var key = meshName + "_" + matName;
            if (lodResMap.ContainsKey(key))
            {
                return lodResMap[key];
            }

            var lodRes = new GPUStaticRenderLODRes();
#if UNITY_EDITOR
            lodRes.mesh = mesh;
            lodRes.material = mat;
#endif
            lodRes.meshName = meshName;
            lodRes.matName = matName;

            var lodResId = lodList.Count;
            lodList.Add(lodRes);
            lodResMap.Add(key, lodResId);
            return lodResId;
        }

        #endregion

    }
}
